"""
Generate test mesh files for Smart Mesh Optimizer
"""
import os
import numpy as np
try:
    import vtk
except ImportError:
    print("VTK is not installed. Please install it with 'pip install vtk'")
    exit(1)

def generate_cube(output_file, size=1.0, quality='good'):
    """
    Generate a cube mesh with specified quality
    
    Parameters:
    -----------
    output_file : str
        Output file path
    size : float
        Size of the cube
    quality : str
        Quality of the mesh ('good', 'medium', 'poor')
    """
    # Create a cube source
    cube = vtk.vtkCubeSource()
    cube.SetXLength(size)
    cube.SetYLength(size)
    cube.SetZLength(size)
    cube.Update()
    
    # Triangulate the cube
    triangulate = vtk.vtkTriangleFilter()
    triangulate.SetInputData(cube.GetOutput())
    triangulate.Update()
    
    mesh = triangulate.GetOutput()
    
    # Deform the mesh based on quality
    if quality != 'good':
        # Get points
        points = mesh.GetPoints()
        n_points = points.GetNumberOfPoints()
        
        # Apply random deformation to interior points
        for i in range(n_points):
            point = list(points.GetPoint(i))
            
            # Skip vertices of the cube
            is_vertex = sum([abs(abs(coord) - size/2) < 0.01 for coord in point]) >= 3
            if is_vertex:
                continue
            
            # Apply deformation
            if quality == 'medium':
                noise = np.random.normal(0, 0.05, 3)
            else:  # poor
                noise = np.random.normal(0, 0.15, 3)
            
            # Ensure points stay within the cube
            for j in range(3):
                point[j] += noise[j]
                point[j] = max(min(point[j], size/2), -size/2)
            
            points.SetPoint(i, point)
        
        mesh.Modified()
    
    # Write to STL file
    writer = vtk.vtkSTLWriter()
    writer.SetFileName(output_file)
    writer.SetInputData(mesh)
    writer.Write()
    
    print(f"Mesh generated: {output_file}")

def generate_sphere(output_file, radius=1.0, resolution=20, quality='good'):
    """
    Generate a sphere mesh with specified quality
    
    Parameters:
    -----------
    output_file : str
        Output file path
    radius : float
        Radius of the sphere
    resolution : int
        Resolution of the sphere
    quality : str
        Quality of the mesh ('good', 'medium', 'poor')
    """
    # Create a sphere source
    sphere = vtk.vtkSphereSource()
    sphere.SetRadius(radius)
    sphere.SetThetaResolution(resolution)
    sphere.SetPhiResolution(resolution)
    sphere.Update()
    
    mesh = sphere.GetOutput()
    
    # Deform the mesh based on quality
    if quality != 'good':
        # Get points
        points = mesh.GetPoints()
        n_points = points.GetNumberOfPoints()
        
        # Apply random deformation
        for i in range(n_points):
            point = list(points.GetPoint(i))
            
            # Apply deformation
            if quality == 'medium':
                noise = np.random.normal(0, 0.05, 3)
            else:  # poor
                noise = np.random.normal(0, 0.15, 3)
            
            # Add noise
            for j in range(3):
                point[j] += noise[j]
            
            points.SetPoint(i, point)
        
        mesh.Modified()
    
    # Write to STL file
    writer = vtk.vtkSTLWriter()
    writer.SetFileName(output_file)
    writer.SetInputData(mesh)
    writer.Write()
    
    print(f"Mesh generated: {output_file}")

def main():
    """Main function"""
    # Create output directory
    output_dir = os.path.dirname(os.path.abspath(__file__))
    
    # Generate cube meshes
    generate_cube(os.path.join(output_dir, 'cube_good.stl'), quality='good')
    generate_cube(os.path.join(output_dir, 'cube_medium.stl'), quality='medium')
    generate_cube(os.path.join(output_dir, 'cube_poor.stl'), quality='poor')
    
    # Generate sphere meshes
    generate_sphere(os.path.join(output_dir, 'sphere_good.stl'), quality='good')
    generate_sphere(os.path.join(output_dir, 'sphere_medium.stl'), quality='medium')
    generate_sphere(os.path.join(output_dir, 'sphere_poor.stl'), quality='poor')

if __name__ == "__main__":
    main() 